package org.apache.hadoop.mapred;

import java.io.IOException;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
/*import java.util.Iterator;
import java.util.Set;
import java.util.HashMap;
import java.util.Vector;*/
import java.util.*;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import edu.ncsu.csc.streamingmr.FixedLengthTextInputStreamFormat;
import edu.ncsu.csc.streamingmr.InputStreamFormat;
import edu.ncsu.csc.streamingmr.StreamBatch;


public class JobDispatcher extends Thread{

	private static final Log LOG = LogFactory.getLog("org.apache.hadoop.mapred.JobDispatcher");
	
	public int internalMapPort;
	public int externalPort;
	public int internalReducePort;
	private final int byteBufferLength = 8192; 
	
	private int externalBytesRead = 0;
	private int totalBytesWriten = 0;
	
	private int totalMaps = 0;
	
	private long time = -1;
	
	FixedLengthTextInputStreamFormat inputFormat;   
	List<StreamBatch> batchList = new ArrayList<StreamBatch>();
	
	List<InetSocketAddress> address = new ArrayList<InetSocketAddress>();
	String strAddress = "";
	
	public String generateAddressString(List<InetSocketAddress> addrList){
		String rtn = "";
		for(InetSocketAddress addr : addrList){
			rtn+= addr.toString() + "|";
		}
		return rtn.substring(0, rtn.length()-1);
	}	

	public JobDispatcher(JobConf job){
		inputFormat=new FixedLengthTextInputStreamFormat(job);   
		internalMapPort = job.getInt("mapred.job.mapper.stream.port", 3001);
		internalReducePort = job.getInt("mapred.job.reducer.stream.port", 3002);
		externalPort = job.getInt("mapred.jobdispatcher.stream.externalport", 3003);
	}

	public void run() {

		System.out.println("Listening for connections on port " + internalMapPort + "\n");

		ServerSocketChannel serverChannel, serverChannel2, serverChannel3;
		Selector selector;

		try {
			serverChannel = ServerSocketChannel.open();
			ServerSocket ss = serverChannel.socket();
			InetSocketAddress address = new InetSocketAddress(internalMapPort);
			ss.bind(address);
			serverChannel.configureBlocking(false);

			selector = Selector.open();
			SelectionKey listenerKey1 = serverChannel.register(selector, SelectionKey.OP_ACCEPT);
			String strListener1 = "1";
			listenerKey1.attach(strListener1);

			
			//open another server channel for the external data
			serverChannel2 = ServerSocketChannel.open();
			ServerSocket ss2 = serverChannel2.socket();
			InetSocketAddress address2 = new InetSocketAddress(externalPort);
			ss2.bind(address2);
			serverChannel2.configureBlocking(false);

			SelectionKey listenerKey2 = serverChannel2.register(selector, SelectionKey.OP_ACCEPT);
			String strListener2 = "2";
			listenerKey2.attach(strListener2);
			
			//open another server channel for the internal reduce control
			serverChannel3 = ServerSocketChannel.open();
			ServerSocket ss3 = serverChannel3.socket();
			InetSocketAddress address3 = new InetSocketAddress(internalReducePort);
			ss3.bind(address3);
			serverChannel3.configureBlocking(false);

			SelectionKey listenerKey3 = serverChannel3.register(selector, SelectionKey.OP_ACCEPT);
			String strListener3 = "3";
			listenerKey3.attach(strListener3);


		} catch (IOException ex) {
			ex.printStackTrace();
			return;
		}

		while (true) {

			try {
				selector.select();
			} catch (IOException ex) {
				ex.printStackTrace();
				break;
			}

			Set readyKeys = selector.selectedKeys();
			Iterator iterator = readyKeys.iterator();

			while (iterator.hasNext()) {

				SelectionKey key = (SelectionKey) iterator.next();
				iterator.remove();

				try {
					if (key.isAcceptable()) {
						ServerSocketChannel server = (ServerSocketChannel) key.channel();
						SocketChannel client = server.accept();
						client.configureBlocking(false);
						SelectionKey key2 = client.register(selector, SelectionKey.OP_READ);

						// if this is a connection from MapTask
						if( ((String)key.attachment()).equals("1")  ){
							InternalMapEventHandler internalMapHandler = new InternalMapEventHandler();
							key2.attach(internalMapHandler);
							totalMaps ++;
							/*Socket socket = client.socket();
							InetSocketAddress incomingAddr = new InetSocketAddress(socket.getInetAddress(), socket.getPort());
							address.add(incomingAddr);
							strAddress = generateAddressString(address);
							System.out.print(strAddress+"\n");*/
							//internalMapHandler.setType(1);
						} else if (((String)key.attachment()).equals("2")) { // if this is a connection from external data source
							ExternalEventHandler externalHandler = new ExternalEventHandler();
							key2.attach(externalHandler);	
							//externalHandler.setType(2);

						} else { // if this is a connection from ReduceTask
							InternalReduceEventHandler internalReduceHandler = new InternalReduceEventHandler();
							key2.attach(internalReduceHandler);	
							//internalReduceHandler.setType(2);
						}
					}
					else if (key.isReadable()) {
						SocketChannel client = (SocketChannel) key.channel();
						EventHandler handler = (EventHandler) key.attachment();
						handler.handleRead(client, key);

					} else if (key.isWritable()) {
						SocketChannel client = (SocketChannel) key.channel();
						EventHandler handler = (EventHandler) key.attachment();
						handler.handleWrite(client, key);
					}


				} catch (IOException ex) {
					key.cancel();
					try {
						key.channel().close();
					}
					catch (IOException cex) {
					}
				}
				
			} // end of while(iterator.hasNext())
			
		} // end of while(true)
	} // end of main()


	class InternalReduceEventHandler extends EventHandler{

		private ByteBuffer outputBuffer = null;
		

		public void handleRead(SocketChannel channel, SelectionKey key){

			try{
				//System.out.print("inside internal event handler\n");
				channel.read(buffer);
				buffer.flip();
				byte[] bytearray = new byte[buffer.remaining()];
				buffer.get(bytearray);
				String s = new String(bytearray);                                                         
				//System.out.print("Reduce: "+s+"\n");
				//System.out.print(channel.socket().getInetAddress().getHostName()+"\n");

				buffer.clear();

				key.interestOps(SelectionKey.OP_WRITE);

			} catch(IOException ioe) {
				LOG.info("Read IOException in InternalEventHandler");

			}
		}

		public void handleWrite(SocketChannel channel, SelectionKey key){
			try{
				if(outputBuffer == null || !outputBuffer.hasRemaining() ){
					if(strAddress != null && strAddress.length()>0){
						byte[] byteAddr = strAddress.getBytes();
						outputBuffer = ByteBuffer.wrap(byteAddr);
					}else{
						channel.close();
						return;
					}
				}
				channel.write(outputBuffer);
				if(!outputBuffer.hasRemaining()){
					channel.close();
				}
			} catch (IOException ioe) {
				LOG.info("Write IOException in InternalEventHandler");
			}
		}
	}
	
	class StreamBatchNode{
		StreamBatch streamBatch;
		StreamBatchNode next;
		ByteBuffer buffer;
		
		public StreamBatchNode(StreamBatch streamBatch){
			this.streamBatch = streamBatch;
			this.buffer = ByteBuffer.wrap(this.streamBatch.getByteArray());
			next = null;
		}
		
		public ByteBuffer getBuffer(){
			return this.buffer;
		}
	}
	
	class OutputBatchList{
		StreamBatchNode first = null;
		StreamBatchNode last = null;
		int count = 0;
		
		public void addTail(StreamBatch streamBatch){
			if(count == 0){
				first = last = new StreamBatchNode(streamBatch);
				count++;
			} else {
				last.next=new StreamBatchNode(streamBatch);
				last=last.next;
				count++;
			}
		}
		
		public ByteBuffer getHead(){
			return first.getBuffer();
		}
		
		public void removeHead(){
			if(count==1){
				first = last =null;
				count--;
			} else if (count>1){
				first = first.next;
				count--;
			} else {
				first = last =null;
				count = 0;
			}
		}
		
		public boolean isEmpty(){
			return (count==0? true : false);
		}
	}


	class InternalMapEventHandler extends EventHandler{

		private ByteBuffer outputBuffer = null;
		OutputBatchList outList = new OutputBatchList();
		String s;
		private final int maxQuota = 100;
		//private long memoryLimit = 10000000; //50MB

		public void handleRead(SocketChannel channel, SelectionKey key){

			try{
				//System.out.print("inside internal event handler\n");
				channel.read(buffer);
				buffer.flip();
				byte[] bytearray = new byte[buffer.remaining()];
				buffer.get(bytearray);
				s = new String(bytearray);                                                         
				//System.out.print("Map: "+s+"\n");
				
				String port = s.substring(0, 4);				
				InetSocketAddress addr = new InetSocketAddress(channel.socket().getInetAddress(), Integer.parseInt(port));
				boolean found = false;
				for(InetSocketAddress a : address){
					if(a.equals(addr)){
						found = true;
						break;
					}
				}
				if(!found){
					address.add(addr);
					strAddress = generateAddressString(address);
					//System.out.print(address.size()+"\n");
				}
				
				buffer.clear();

				key.interestOps(SelectionKey.OP_WRITE);

			} catch(IOException ioe) {
				LOG.info("Read IOException in InternalEventHandler");

			}
		}
		
		public void handleWrite(SocketChannel channel, SelectionKey key){
			try{
				//LOG.info("write handler called");
				if(outList.isEmpty()){
					
					if(batchList.size()>0){
						//System.out.print("BEFORE: batchList.size()="+batchList.size()+"\n");
						int currentSize = batchList.size();
						int perMapQuota = Math.max(1, currentSize/totalMaps);
						perMapQuota = Math.min(perMapQuota, maxQuota);
						//int totalLength = 0;
						//int upperLimit = perMapQuota;
						//for(int j=0; j<perMapQuota; j++){
							//totalLength += ((StreamBatch)batchList.get(j)).getBatchSize();
							/*if(totalLength >= memoryLimit){
								totalLength -= ((StreamBatch)batchList.get(j)).getBatchSize();
								upperLimit = j;
								break;
							}*/
							//System.out.println("testfji before writing, batch="+new String(batch));
						//}
						//System.out.println("  totalLength="+totalLength);
						//byte[] batch = new byte[totalLength];
						//int pos=0;
						//System.out.print("perMapQuota = "+perMapQuota+"\n");
						for(int j=0; j<perMapQuota; j++){
							StreamBatch temp = (StreamBatch)batchList.remove(0);
							outList.addTail(temp);
							//System.arraycopy(temp.getByteArray(), 0, batch, pos, temp.getBatchSize());
							//pos+=temp.getBatchSize();
						}
						//System.out.print("AFTER: batchList.size()="+batchList.size()+"\n");
						//System.out.println("here");
						//outputBuffer = ByteBuffer.wrap(batch);
						/*if(batchList.size()==0){
							System.out.println("Batch size becomes 0 at: "+System.currentTimeMillis());
						}*/
					}
				}
				//System.out.print("start writing\n");
				while(!outList.isEmpty()){
					outputBuffer = outList.getHead();
					channel.write(outputBuffer);
					if(!outputBuffer.hasRemaining()){
						outList.removeHead();
					}else{
						break;
					}
					//System.out.println("map channel isblocking() = "+channel.isBlocking());
					//totalBytesWriten += channel.write(outputBuffer);
					//System.out.print("totalBytesWriten: -----------------------"+totalBytesWriten+"\n");
					//LOG.info("data was written\n");
				}
				//System.out.print("finish writing\n\n");
			} catch (IOException ioe) {
				LOG.info("Write IOException in InternalEventHandler");
			}
		}
	}




	class ExternalEventHandler extends EventHandler{
		
		/*int leftoverLength = 40;
		byte[] leftoverBuffer;
		int leftoverBufferLength = 0;*/
			
		
		public void handleRead(SocketChannel channel, SelectionKey key){
			try{
				//System.out.println("external channel isblocking() = "+channel.isBlocking());
				if(time==-1){
					time=System.currentTimeMillis();
					System.out.println("STARTING TIME: "+time);
				}
				
				//System.out.print("\nread called\n");
				int count = channel.read(buffer);
				//System.out.print("\nread finished\n");
				if(count > 0){
					externalBytesRead += count;
					//System.out.print("externalBytesRead: ----------------------------------------- " + externalBytesRead + "\n");
					buffer.flip();
					byte[] rawData = new byte[buffer.remaining()];
					buffer.get(rawData);
					StreamBatch[] batches = inputFormat.getStreamBatch(rawData);
					for(StreamBatch batch : batches){
						batchList.add(batch);
					}
					buffer.clear();
				} else if (count == -1) {
					/*System.out.print("count = "+count+"\n");
					channel.close();*/
				}
				//System.out.print("total "+count+" bytes read\n");
			} catch (IOException ioe) {
				LOG.info("Read IOException in ExternalEventHandler");
			}
		}

		public void handleWrite(SocketChannel channel, SelectionKey key){}

	}
	

	abstract class EventHandler {
			
			protected int type;
			protected ByteBuffer buffer;

			public EventHandler(){
				buffer = ByteBuffer.allocate(byteBufferLength);
			}

			public ByteBuffer getBuffer(){
				return this.buffer;
			}


			public int getType(){
				return this.type;
			}

			//1--internal connection
			//2--external connection
			public void setType(int type){
				this.type = type;
			}

			public abstract void handleRead(SocketChannel channel, SelectionKey key);
			public abstract void handleWrite(SocketChannel channel, SelectionKey key);

	}

	
	public static void main(String args[]){
		//JobDispatcher jd = new JobDispatcher();
		//jd.start();
	}
}
